import { logger } from 'logging/logger';
import {modelChanged} from 'features/parameters/store/generationSlice';
import { zMainOrOnnxModel } from 'features/parameters/types/parameterSchemas';

import {
  mainModelsAdapter,
  modelsApi,
} from 'services/api/endpoints/models';
import { TypeGuardFor } from 'services/api/types';
import { startAppListening } from '..';

export const addModelsLoadedListener = () => {
  startAppListening({
    predicate: (
      action
    ): action is TypeGuardFor<
      typeof modelsApi.endpoints.getMainModels.matchFulfilled
    > =>
      modelsApi.endpoints.getMainModels.matchFulfilled(action) &&
      !action.meta.arg.originalArgs.includes('sdxl-refiner'),
    effect: async (action, { getState, dispatch }) => {
      // models loaded, we need to ensure the selected model is available and if not, select the first one
      const log = logger('models');
      log.info(
        { models: action.payload.entities },
        `Main models loaded (${action.payload.ids.length})`
      );

      const currentModel = getState().generation.model;
      const models = mainModelsAdapter.getSelectors().selectAll(action.payload);

      if (models.length === 0) {
        // No models loaded at all
        dispatch(modelChanged(null));
        return;
      }

      const isCurrentModelAvailable = currentModel
        ? models.some(
            (m) =>
              m.model_name === currentModel.model_name &&
              m.base_model === currentModel.base_model &&
              m.model_type === currentModel.model_type
          )
        : false;

      if (isCurrentModelAvailable) {
        return;
      }

      const result = zMainOrOnnxModel.safeParse(models[0]);

      if (!result.success) {
        log.error(
          { error: result.error.format() },
          'Failed to parse main model'
        );
        return;
      }

      dispatch(modelChanged(result.data));
    },
  });
};
